import numpy as np


class StockDataset:

    def __init__(self, data):
        self.data = data
        self.per = 0.085

    def split(self):
        test_set_size = int(np.round(self.per * self.data.shape[0]))
        train_set_size = self.data.shape[0] - test_set_size
        x_train = self.data.iloc[:train_set_size, 1:-1]

        x_train["open"] = x_train["open"] - x_train["zl1"]
        x_train["close"] = x_train["close"] - x_train["zl1"]
        x_train["high"] = x_train["high"] - x_train["zl1"]
        x_train["low"] = x_train["low"] - x_train["zl1"]
        x_train = x_train[["open", "close", "high", "low", "change"]]

        y_train = self.data.iloc[:train_set_size, -1]

        x_test = self.data.iloc[train_set_size:, 1:-1]

        x_test["open"] = x_test["open"] - x_test["zl1"]
        x_test["close"] = x_test["close"] - x_test["zl1"]
        x_test["high"] = x_test["high"] - x_test["zl1"]
        x_test["low"] = x_test["low"] - x_test["zl1"]
        x_test = x_test[["open", "close", "high", "low", "change"]]

        y_test = self.data.iloc[train_set_size:, -1]

        return x_train, y_train, x_test, y_test
